package main

import (
	"flag"
	"fmt"
	"net"
	"os"
)

func main() {
	/*
	   定义变量接收控制台参数
	*/

	//连接的ip地址
	var ip string
	//远端端口
	var rp int
	//本地端口
	var lp int

	// StringVar用指定的名称、控制台参数项目、默认值、使用信息注册一个string类型flag，并将flag的值保存到p指向的变量
	flag.StringVar(&ip, "ip", "", "转发的远程ip")
	flag.IntVar(&rp, "rp", 0, "转发的远程端口")
	flag.IntVar(&lp, "lp", 0, "本地监听的端口")

	// 从arguments中解析注册的flag。必须在所有flag都注册好而未访问其值时执行。未注册却使用flag -help时，会返回ErrHelp。
	flag.Parse()

	// 打印
	//fmt.Printf("ip=%v,remotePort=%v,localPort=%v\n", ip, rp, lp)

	//localPort, err1 := strconv.Atoi(lp)
	//remotePort, err2 := strconv.Atoi(rp)
	//if err1 != nil {
	//	fmt.Printf("本地监听端口输入错误:%v\n", err1)
	//	return
	//}
	//
	//if err2 != nil {
	//	fmt.Printf("远程端口输入错误:%v\n", err2)
	//	return
	//}
	if flag.NFlag() != 3 {
		println("params error! 参数输入错误")
		flag.Usage()
		os.Exit(0)
	}
	proxyStart(lp, rp, ip)

	//proxyStart(7991, 22, "149.129.71.43")

}

// Start a proxy server listen on fromport
// this proxy will then forward all request from fromport to toport
//
// Notice: a service must has been started on toport
func proxyStart(localport, remoteport int, remoteaddr string) {
	proxyaddr := fmt.Sprintf(":%d", localport)
	proxylistener, err := net.Listen("tcp", proxyaddr)
	if err != nil {
		fmt.Println("Unable to listen on: %s, error: %s\n", proxyaddr, err.Error())
		os.Exit(1)
	}
	defer proxylistener.Close()
	fmt.Printf("listing port :%v, forward to %v:%v starting...", localport, remoteaddr, remoteport)
	for {
		proxyconn, err := proxylistener.Accept()
		if err != nil {
			fmt.Printf("Unable to accept a request, error: %s\n", err.Error())
			continue
		}

		// Read a header firstly in case you could have opportunity to check request
		// whether to decline or proceed the request
		buffer := make([]byte, 1024)
		n, err := proxyconn.Read(buffer)
		if err != nil {
			fmt.Printf("Unable to read from input, error: %s\n", err.Error())
			continue
		}

		// TODO
		// Your choice to make decision based on request header

		//targetaddr := fmt.Sprintf("localhost:%d", toport)
		targetaddr := fmt.Sprintf("%s:%d", remoteaddr, remoteport)
		targetconn, err := net.Dial("tcp", targetaddr)
		if err != nil {
			fmt.Println("Unable to connect to: %s, error: %s\n", targetaddr, err.Error())
			proxyconn.Close()
			continue
		}

		n, err = targetconn.Write(buffer[:n])
		if err != nil {
			fmt.Printf("Unable to write to output, error: %s\n", err.Error())
			proxyconn.Close()
			targetconn.Close()
			continue
		}

		go proxyRequest(proxyconn, targetconn)
		go proxyRequest(targetconn, proxyconn)
	}
}

// Forward all requests from r to w
func proxyRequest(r net.Conn, w net.Conn) {
	defer r.Close()
	defer w.Close()

	var buffer = make([]byte, 4096000)
	for {
		n, err := r.Read(buffer)
		if err != nil {
			fmt.Printf("Unable to read from input, error: %s\n", err.Error())
			break
		}

		n, err = w.Write(buffer[:n])
		if err != nil {
			fmt.Printf("Unable to write to output, error: %s\n", err.Error())
			break
		}
	}
}
